from langgraph.graph import StateGraph, END

from app.agent.nodes import load_data, preprocess_data, generate_factors, backtest_factor, evaluate_factors, \
    fx_feature_engineering, bond_feature_engineering, metal_feature_engineering
from app.agent.quant_params import QuantParams
from app.agent.state import QuantState


def init():
    # 初始化图
    graph = StateGraph(QuantState)

    # --------------- 构建工作流 ---------------
    graph.add_node("load_data", load_data)
    graph.add_node("preprocess", preprocess_data)
    graph.add_node("fx_engineer", fx_feature_engineering)
    graph.add_node("bond_engineer", bond_feature_engineering)
    graph.add_node("metal_engineer", metal_feature_engineering)
    graph.add_node("generate_factors", generate_factors)
    graph.add_node("backtest", backtest_factor)
    graph.add_node("evaluate", evaluate_factors)

    # 定义边（线性流程+条件分支）
    graph.set_entry_point("load_data")
    graph.add_edge("load_data", "preprocess")
    graph.add_edge("fx_engineer", "generate_factors")
    graph.add_edge("bond_engineer", "generate_factors")
    graph.add_edge("metal_engineer", "generate_factors")
    graph.add_edge("generate_factors", "backtest")
    graph.add_edge("backtest", "evaluate")

    # 可选：添加循环优化分支
    def should_optimize(state: QuantState) -> str:
        print("===========should_optimize===========")
        print(state)
        if len(state["selected_factors"]) < 1:  # 如果有效因子不足，重新生成
            return "continue"
        else:
            return END

    graph.add_conditional_edges(
        "evaluate",
        should_optimize,
        {
            # If `tools`, then we call the tool node.
            "continue": "generate_factors",
            # Otherwise we finish.
            "__end__": END,
        },
    )

    # 条件路由
    def route_by_asset(state: QuantState):
        if state["params"]["asset_class"] == "外汇":
            return "fx_engineer"
        elif state["params"]["asset_class"] == "债券":
            return "metal_engineer"
        else:
            return "bond_engineer"

    graph.add_conditional_edges(
        "preprocess",
        route_by_asset,
        {"fx_engineer": "fx_engineer", "bond_engineer": "bond_engineer"}
    )

    # 编译并运行
    app = graph.compile()
    return app


def start(app):
    params = QuantParams(
        asset_class="国债",
        bond_duration="中端(3-7Y)",
        time_frame="日线",
        max_drawdown=0.1,
        data_sources=["宏观经济指标", "央行政策文本"],
        credit_rating="AAA"
    )
    state = QuantState(params=params, features={}, raw_data=None, processed_data=None,
                       candidate_factors=[], backtest_results={}, selected_factors=[])

    result = app.invoke(state)
    # socket_handler.send_message("[agent][执行结束]"+json.dumps(result["selected_factors"]))


if __name__ == '__main__':
    app = init()
    start(app)
